use crate::config::SecurityConfig;
use crate::errors::{AppError, AppResult};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use uuid::Uuid;

/// Security manager for handling authentication, authorization, and rate limiting
pub struct SecurityManager {
    /// Configuration
    config: SecurityConfig,
    /// API key storage
    api_keys: Arc<RwLock<HashMap<String, ApiKeyInfo>>>,
    /// Rate limiter
    rate_limiter: Arc<RateLimiter>,
    /// Content filter
    content_filter: Arc<ContentFilter>,
    /// Session manager
    session_manager: Arc<SessionManager>,
    /// Audit logger
    audit_logger: Arc<AuditLogger>,
}

/// API key information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ApiKeyInfo {
    /// API key value (hashed)
    pub key_hash: String,
    /// Key name/identifier
    pub name: String,
    /// Key owner
    pub owner: String,
    /// Key permissions
    pub permissions: HashSet<String>,
    /// Key expiration
    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
    /// Key creation time
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// Last used time
    pub last_used: Option<chrono::DateTime<chrono::Utc>>,
    /// Whether the key is active
    pub is_active: bool,
    /// Rate limit override
    pub rate_limit_override: Option<u32>,
}

/// Rate limiter implementation
pub struct RateLimiter {
    /// Request records by client ID
    records: Arc<RwLock<HashMap<String, RequestRecord>>>,
    /// Default rate limit (requests per minute)
    default_limit: u32,
    /// Cleanup interval
    cleanup_interval: Duration,
}

/// Request record for rate limiting
#[derive(Debug, Clone)]
pub struct RequestRecord {
    /// Request timestamps
    requests: Vec<Instant>,
    /// Current count
    count: u32,
    /// Rate limit
    limit: u32,
}

/// Content filter for detecting malicious content
pub struct ContentFilter {
    /// Blocked keywords
    blocked_keywords: HashSet<String>,
    /// Allowed file extensions
    allowed_extensions: HashSet<String>,
    /// Maximum content length
    max_content_length: usize,
    /// Enable content scanning
    enable_scanning: bool,
}

/// Session information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionInfo {
    /// Session ID
    pub session_id: String,
    /// User ID
    pub user_id: String,
    /// Client IP
    pub client_ip: String,
    /// User agent
    pub user_agent: String,
    /// Session creation time
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// Last activity time
    pub last_activity: chrono::DateTime<chrono::Utc>,
    /// Session expiration
    pub expires_at: chrono::DateTime<chrono::Utc>,
    /// Session permissions
    pub permissions: HashSet<String>,
    /// Whether session is active
    pub is_active: bool,
}

/// Session manager
pub struct SessionManager {
    /// Active sessions
    sessions: Arc<RwLock<HashMap<String, SessionInfo>>>,
    /// Session timeout
    session_timeout: Duration,
}

/// Audit log entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditLogEntry {
    /// Entry ID
    pub id: String,
    /// Timestamp
    pub timestamp: chrono::DateTime<chrono::Utc>,
    /// User ID
    pub user_id: Option<String>,
    /// Session ID
    pub session_id: Option<String>,
    /// Action performed
    pub action: String,
    /// Resource affected
    pub resource: Option<String>,
    /// Result
    pub result: AuditResult,
    /// Client IP
    pub client_ip: Option<String>,
    /// Additional details
    pub details: HashMap<String, serde_json::Value>,
}

/// Audit result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AuditResult {
    Success,
    Failure,
    Blocked,
    Warning,
}

/// Audit logger for security events
pub struct AuditLogger {
    /// Audit log entries
    logs: Arc<RwLock<Vec<AuditLogEntry>>>,
}

impl SecurityManager {
    /// Create a new security manager
    pub fn new(config: SecurityConfig) -> Self {
        let rate_limiter = Arc::new(RateLimiter::new(config.max_requests_per_minute));
        let content_filter = Arc::new(ContentFilter::new(&config));
        let session_manager = Arc::new(SessionManager::new());
        let audit_logger = Arc::new(AuditLogger::new());

        Self {
            config,
            api_keys: Arc::new(RwLock::new(HashMap::new())),
            rate_limiter,
            content_filter,
            session_manager,
            audit_logger,
        }
    }

    /// Initialize security manager
    pub async fn initialize(&self) -> AppResult<()> {
        // Load default API keys from environment if configured
        if let Ok(default_key) = std::env::var("RIG_DEFAULT_API_KEY") {
            self.create_api_key(
                "default".to_string(),
                "system".to_string(),
                default_key,
                vec!["*".to_string()],
                None,
            )
            .await?;
        }

        // Start background cleanup tasks
        self.start_cleanup_tasks().await?;

        Ok(())
    }

    /// Create a new API key
    pub async fn create_api_key(
        &self,
        name: String,
        owner: String,
        key_value: String,
        permissions: Vec<String>,
        expires_at: Option<chrono::DateTime<chrono::Utc>>,
    ) -> AppResult<String> {
        let key_hash = self.hash_api_key(&key_value)?;
        let key_id = Uuid::new_v4().to_string();

        let key_info = ApiKeyInfo {
            key_hash,
            name: name.clone(),
            owner,
            permissions: permissions.clone().into_iter().collect(),
            expires_at,
            created_at: chrono::Utc::now(),
            last_used: None,
            is_active: true,
            rate_limit_override: None,
        };

        let mut keys = self.api_keys.write().await;
        keys.insert(key_id.clone(), key_info);

        // Log API key creation
        self.audit_logger
            .log(
                None,
                None,
                "api_key_created".to_string(),
                Some(format!("API key '{}' created", name)),
                AuditResult::Success,
                None,
                Some({
                    let mut details = HashMap::new();
                    details.insert("key_name".to_string(), name.into());
                    details.insert("permissions".to_string(), serde_json::json!(permissions));
                    details
                }),
            )
            .await;

        Ok(key_id)
    }

    /// Validate an API key
    pub async fn validate_api_key(&self, key_value: &str) -> AppResult<(String, ApiKeyInfo)> {
        if !self.config.validate_api_keys {
            // Skip validation if disabled
            return Ok(("anonymous".to_string(), ApiKeyInfo::default()));
        }

        let key_hash = self.hash_api_key(key_value)?;
        let keys = self.api_keys.read().await;

        // Find matching key and clone the data we need
        let (key_id, key_info) = keys
            .iter()
            .find(|(_, info)| info.key_hash == key_hash && info.is_active)
            .ok_or_else(|| AppError::auth("Invalid API key"))?;

        // Check expiration
        if let Some(expires_at) = key_info.expires_at {
            if chrono::Utc::now() > expires_at {
                return Err(AppError::auth("API key expired"));
            }
        }

        // Clone the data we need before dropping the read lock
        let key_id = key_id.clone();
        let key_info = key_info.clone();

        // Update last used time
        drop(keys);
        let mut keys = self.api_keys.write().await;
        if let Some(key_info) = keys.get_mut(&key_id) {
            key_info.last_used = Some(chrono::Utc::now());
        }

        Ok((key_id, key_info))
    }

    /// Check rate limit for a client
    pub async fn check_rate_limit(
        &self,
        client_id: &str,
        custom_limit: Option<u32>,
    ) -> AppResult<()> {
        if !self.config.enable_rate_limiting {
            return Ok(());
        }

        let limit = custom_limit.unwrap_or(self.config.max_requests_per_minute);
        self.rate_limiter.check_rate_limit(client_id, limit).await
    }

    /// Filter content for security
    pub async fn filter_content(&self, content: &str, content_type: &str) -> AppResult<()> {
        if !self.config.enable_content_filtering {
            return Ok(());
        }

        self.content_filter
            .filter_content(content, content_type)
            .await
    }

    /// Validate file upload
    pub async fn validate_file_upload(&self, filename: &str, content: &[u8]) -> AppResult<()> {
        // Check file extension
        let extension = filename.split('.').last().unwrap_or("");
        if !self
            .config
            .allowed_file_extensions
            .contains(&extension.to_string())
        {
            return Err(AppError::validation(format!(
                "File extension '.{}' is not allowed",
                extension
            )));
        }

        // Check file size (max 10MB)
        if content.len() > 10 * 1024 * 1024 {
            return Err(AppError::validation("File size exceeds limit (10MB)"));
        }

        // Scan content if enabled
        if self.config.enable_content_filtering {
            let content_str = String::from_utf8_lossy(content);
            self.filter_content(&content_str, "file").await?;
        }

        Ok(())
    }

    /// Create a new session
    pub async fn create_session(
        &self,
        user_id: String,
        client_ip: String,
        user_agent: String,
        permissions: HashSet<String>,
    ) -> AppResult<String> {
        self.session_manager
            .create_session(user_id, client_ip, user_agent, permissions)
            .await
    }

    /// Validate a session
    pub async fn validate_session(&self, session_id: &str) -> AppResult<SessionInfo> {
        self.session_manager.validate_session(session_id).await
    }

    /// Check permission for a session
    pub async fn check_permission(&self, session_id: &str, permission: &str) -> AppResult<()> {
        let session = self.validate_session(session_id).await?;

        if !session.permissions.contains(permission) && !session.permissions.contains("*") {
            return Err(AppError::permission(format!(
                "Permission '{}' denied",
                permission
            )));
        }

        Ok(())
    }

    /// Log an audit event
    pub async fn log_audit_event(
        &self,
        user_id: Option<String>,
        session_id: Option<String>,
        action: String,
        resource: Option<String>,
        result: AuditResult,
        client_ip: Option<String>,
        details: Option<HashMap<String, serde_json::Value>>,
    ) {
        self.audit_logger
            .log(
                user_id, session_id, action, resource, result, client_ip, details,
            )
            .await;
    }

    /// Get audit logs
    pub async fn get_audit_logs(&self, limit: usize, offset: usize) -> Vec<AuditLogEntry> {
        self.audit_logger.get_logs(limit, offset).await
    }

    /// Start background cleanup tasks
    async fn start_cleanup_tasks(&self) -> AppResult<()> {
        let rate_limiter = self.rate_limiter.clone();
        let session_manager = self.session_manager.clone();

        // Start rate limiter cleanup
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            loop {
                interval.tick().await;
                if let Err(e) = rate_limiter.cleanup_expired_records().await {
                    eprintln!("Rate limiter cleanup failed: {}", e);
                }
            }
        });

        // Start session cleanup
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(300)); // 5 minutes
            loop {
                interval.tick().await;
                if let Err(e) = session_manager.cleanup_expired_sessions().await {
                    eprintln!("Session cleanup failed: {}", e);
                }
            }
        });

        Ok(())
    }

    /// Hash API key for storage
    fn hash_api_key(&self, key: &str) -> AppResult<String> {
        use sha2::{Digest, Sha256};
        let mut hasher = Sha256::new();
        hasher.update(key.as_bytes());
        Ok(format!("{:x}", hasher.finalize()))
    }
}

impl RateLimiter {
    /// Create a new rate limiter
    pub fn new(default_limit: u32) -> Self {
        Self {
            records: Arc::new(RwLock::new(HashMap::new())),
            default_limit,
            cleanup_interval: Duration::from_secs(60),
        }
    }

    /// Check if a request is allowed
    pub async fn check_rate_limit(&self, client_id: &str, limit: u32) -> AppResult<()> {
        let mut records = self.records.write().await;
        let now = Instant::now();
        let one_minute_ago = now - Duration::from_secs(60);

        let record = records
            .entry(client_id.to_string())
            .or_insert_with(|| RequestRecord {
                requests: Vec::new(),
                count: 0,
                limit,
            });

        // Update limit if changed
        record.limit = limit;

        // Remove old requests
        record
            .requests
            .retain(|&timestamp| timestamp > one_minute_ago);
        record.count = record.requests.len() as u32;

        // Check if limit exceeded
        if record.count >= record.limit {
            return Err(AppError::rate_limit(format!(
                "Rate limit exceeded ({} requests per minute)",
                record.limit
            )));
        }

        // Add current request
        record.requests.push(now);
        record.count += 1;

        Ok(())
    }

    /// Cleanup expired records
    pub async fn cleanup_expired_records(&self) -> AppResult<()> {
        let mut records = self.records.write().await;
        let now = Instant::now();
        let one_hour_ago = now - Duration::from_secs(3600);

        records.retain(|_, record| {
            record
                .requests
                .iter()
                .any(|&timestamp| timestamp > one_hour_ago)
        });

        Ok(())
    }
}

impl ContentFilter {
    /// Create a new content filter
    pub fn new(config: &SecurityConfig) -> Self {
        let blocked_keywords = [
            "password",
            "secret",
            "token",
            "key",
            "credential",
            "auth",
            "login",
            "signin",
            "admin",
            "root",
            "sudo",
            "exec",
            "eval",
            "system",
            "shell",
            "bash",
            "cmd",
            "powershell",
            "script",
        ]
        .iter()
        .map(|s| s.to_string())
        .collect();

        Self {
            blocked_keywords,
            allowed_extensions: config.allowed_file_extensions.iter().cloned().collect(),
            max_content_length: 10000, // 10KB max content length
            enable_scanning: config.enable_content_filtering,
        }
    }

    /// Filter content for security violations
    pub async fn filter_content(&self, content: &str, content_type: &str) -> AppResult<()> {
        if !self.enable_scanning {
            return Ok(());
        }

        // Check content length
        if content.len() > self.max_content_length {
            return Err(AppError::validation("Content too long"));
        }

        // Check for blocked keywords (case-insensitive)
        let content_lower = content.to_lowercase();
        for keyword in &self.blocked_keywords {
            if content_lower.contains(keyword) {
                return Err(AppError::validation(format!(
                    "Content contains blocked keyword: '{}'",
                    keyword
                )));
            }
        }

        // Additional content type specific checks
        match content_type {
            "prompt" | "query" => {
                // Check for potential injection attempts
                if content.contains("rm -rf") || content.contains("format c:") {
                    return Err(AppError::validation(
                        "Content contains potentially harmful commands",
                    ));
                }
            }
            _ => {}
        }

        Ok(())
    }
}

impl SessionManager {
    /// Create a new session manager
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            session_timeout: Duration::from_secs(3600), // 1 hour timeout
        }
    }

    /// Create a new session
    pub async fn create_session(
        &self,
        user_id: String,
        client_ip: String,
        user_agent: String,
        permissions: HashSet<String>,
    ) -> AppResult<String> {
        let session_id = Uuid::new_v4().to_string();
        let now = chrono::Utc::now();
        let expires_at = now + chrono::Duration::seconds(self.session_timeout.as_secs() as i64);

        let session = SessionInfo {
            session_id: session_id.clone(),
            user_id,
            client_ip,
            user_agent,
            created_at: now,
            last_activity: now,
            expires_at,
            permissions,
            is_active: true,
        };

        let mut sessions = self.sessions.write().await;
        sessions.insert(session_id.clone(), session);

        Ok(session_id)
    }

    /// Validate a session
    pub async fn validate_session(&self, session_id: &str) -> AppResult<SessionInfo> {
        let sessions = self.sessions.read().await;
        let session = sessions
            .get(session_id)
            .ok_or_else(|| AppError::auth("Invalid session ID"))?;

        if !session.is_active {
            return Err(AppError::auth("Session is not active"));
        }

        if chrono::Utc::now() > session.expires_at {
            return Err(AppError::auth("Session has expired"));
        }

        Ok(session.clone())
    }

    /// Update session activity
    pub async fn update_activity(&self, session_id: &str) -> AppResult<()> {
        let mut sessions = self.sessions.write().await;
        if let Some(session) = sessions.get_mut(session_id) {
            session.last_activity = chrono::Utc::now();
            Ok(())
        } else {
            Err(AppError::auth("Session not found"))
        }
    }

    /// Cleanup expired sessions
    pub async fn cleanup_expired_sessions(&self) -> AppResult<()> {
        let mut sessions = self.sessions.write().await;
        let now = chrono::Utc::now();

        sessions.retain(|_, session| session.is_active && session.expires_at > now);

        Ok(())
    }
}

impl AuditLogger {
    /// Create a new audit logger
    pub fn new() -> Self {
        Self {
            logs: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Log an audit event
    pub async fn log(
        &self,
        user_id: Option<String>,
        session_id: Option<String>,
        action: String,
        resource: Option<String>,
        result: AuditResult,
        client_ip: Option<String>,
        details: Option<HashMap<String, serde_json::Value>>,
    ) {
        let entry = AuditLogEntry {
            id: Uuid::new_v4().to_string(),
            timestamp: chrono::Utc::now(),
            user_id,
            session_id,
            action,
            resource,
            result,
            client_ip,
            details: details.unwrap_or_default(),
        };

        let mut logs = self.logs.write().await;
        logs.push(entry);

        // Keep only last 10000 entries
        if logs.len() > 10000 {
            logs.remove(0);
        }
    }

    /// Get audit logs
    pub async fn get_logs(&self, limit: usize, offset: usize) -> Vec<AuditLogEntry> {
        let logs = self.logs.read().await;
        let start_idx = if offset < logs.len() {
            offset
        } else {
            logs.len()
        };
        let end_idx = if start_idx + limit < logs.len() {
            start_idx + limit
        } else {
            logs.len()
        };
        logs[start_idx..end_idx].to_vec()
    }
}

impl Default for ApiKeyInfo {
    fn default() -> Self {
        Self {
            key_hash: String::new(),
            name: "anonymous".to_string(),
            owner: "anonymous".to_string(),
            permissions: ["*".to_string()].into(),
            expires_at: None,
            created_at: chrono::Utc::now(),
            last_used: None,
            is_active: true,
            rate_limit_override: None,
        }
    }
}

// Security middleware implementation
pub struct SecurityMiddleware {
    security_manager: Arc<SecurityManager>,
}

impl SecurityMiddleware {
    /// Create new security middleware
    pub fn new(security_manager: Arc<SecurityManager>) -> Self {
        Self { security_manager }
    }

    /// Middleware function for API requests
    pub async fn authenticate_request(
        &self,
        api_key: Option<&str>,
        client_ip: &str,
    ) -> AppResult<(String, ApiKeyInfo)> {
        match api_key {
            Some(key) => {
                let (key_id, key_info) = self.security_manager.validate_api_key(key).await?;

                // Check rate limit
                self.security_manager
                    .check_rate_limit(&key_id, key_info.rate_limit_override)
                    .await?;

                Ok((key_id, key_info))
            }
            None => {
                // Check anonymous rate limit
                self.security_manager
                    .check_rate_limit(client_ip, None)
                    .await?;

                // Return anonymous access
                Ok(("anonymous".to_string(), ApiKeyInfo::default()))
            }
        }
    }

    /// Validate request content
    pub async fn validate_request_content(
        &self,
        content: &str,
        content_type: &str,
    ) -> AppResult<()> {
        self.security_manager
            .filter_content(content, content_type)
            .await
    }

    /// Log API request
    pub async fn log_request(
        &self,
        key_id: &str,
        action: String,
        resource: Option<String>,
        result: AuditResult,
        client_ip: &str,
        details: Option<HashMap<String, serde_json::Value>>,
    ) {
        self.security_manager
            .log_audit_event(
                Some(key_id.to_string()),
                None,
                action,
                resource,
                result,
                Some(client_ip.to_string()),
                details,
            )
            .await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_api_key_creation() {
        let config = SecurityConfig::default();
        let security = SecurityManager::new(config);

        let key_id = security
            .create_api_key(
                "test_key".to_string(),
                "test_user".to_string(),
                "test_secret_key".to_string(),
                vec!["read".to_string(), "write".to_string()],
                None,
            )
            .await
            .unwrap();

        assert!(!key_id.is_empty());
    }

    #[tokio::test]
    async fn test_rate_limiting() {
        let config = SecurityConfig::default();
        let security = SecurityManager::new(config);

        // Should allow first few requests
        for i in 0..5 {
            let result = security.check_rate_limit("test_client", Some(10)).await;
            assert!(result.is_ok(), "Request {} should be allowed", i);
        }

        // Should eventually block if we exceed limit
        // Note: This test depends on the timing of requests
    }

    #[tokio::test]
    async fn test_content_filtering() {
        let config = SecurityConfig::default();
        let security = SecurityManager::new(config);

        // Should allow safe content
        let result = security.filter_content("Hello world", "text").await;
        assert!(result.is_ok());

        // Should block content with blocked keywords
        let result = security
            .filter_content("Give me your password", "text")
            .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_session_management() {
        let config = SecurityConfig::default();
        let security = SecurityManager::new(config);

        let session_id = security
            .create_session(
                "test_user".to_string(),
                "127.0.0.1".to_string(),
                "test-agent".to_string(),
                ["read".to_string(), "write".to_string()].into(),
            )
            .await
            .unwrap();

        let session = security.validate_session(&session_id).await.unwrap();
        assert_eq!(session.user_id, "test_user");
        assert!(session.permissions.contains("read"));
    }
}
